# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
import random
import torch
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


def random_factor_pair(M):
    # 找出所有正因子
    factors = [i for i in range(1, M+1) if M % i == 0]
    # 随机选择一个因子A
    A = random.choice(factors)
    B = M // A
    return A, B

@GENERATOR_REGISTRY.register("generator_ascend_generate_moe_finalize_routing_v2")
class AscendMoeFinalizeRoutingV2(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        expandedX = case_config.inputs[0]
        expandedRowIdx = case_config.inputs[1]
        x1Optional = case_config.inputs[2]
        x2Optional = case_config.inputs[3]
        biasOptional = case_config.inputs[4]
        scalesOptional = case_config.inputs[5]
        expertIdxOptional = case_config.inputs[6]
        dropPadMode = case_config.inputs[7]
        out = case_config.inputs[8]
        # print(expandedRowIdx.range_values)
        if len(expandedX.shape) == 3: # drop pad 场景 dropPadMode仅支持1或3
            if dropPadMode.range_values == 0:
                dropPadMode.range_values = 1
            elif dropPadMode.range_values == 2:
                dropPadMode.range_values = 3
            E, C, H = expandedX.shape[0], expandedX.shape[1], expandedX.shape[2]
            K = scalesOptional.shape[1]
            if K > E:
                K = E
                scalesOptional.shape[1] = K # K: 为从总的专家E中选出K个专家;  E: expert num，即专家数，E需要大于等于K; 
            Num_Rows = out.shape[0]
            out.shape[1] = H
            expandedRowIdx.shape = [Num_Rows * K] # 要求是一个1D的Tensor,其shape支持（NUM\_ROWS \* K）
            expandedRowIdx.range_values = [-1, E * C - 1] # dropPadMode参数值为1、3时，Tensor中的值取值范围是[-1, E\*C - 1]。
            x1Optional.shape = [Num_Rows, H] # x1Optional shape要求与out的shape一致
            x1Optional.dtype = expandedX.dtype # x1Optional 数据类型要求与expandedX一致
            x2Optional.shape = [Num_Rows, H] # x2Optional shape要求与out的shape一致
            x2Optional.dtype = expandedX.dtype # x2Optional 数据类型要求与expandedX一致
            biasOptional.shape = [E, H] # 限制：其shape支持（E，H）
            biasOptional.dtype = expandedX.dtype # 数据类型要求与expandedX一致
            scalesOptional.shape = [Num_Rows, K] # 限制：其shape支持（NUM\_ROWS，K）
            scalesOptional.dtype = expandedX.dtype # 数据类型要求与expandedX一致
            expertIdxOptional.shape = [Num_Rows, K] # 限制：其shape支持（NUM\_ROWS，K）
            expertIdxOptional.range_values = [0, E -1] # 限制：Tensor中的值取值范围是[0, E-1]
            out.dtype = expandedX.dtype

        elif len(expandedX.shape) == 2: # drop less 场景 dropPadMode仅支持0或2
            if dropPadMode.range_values == 1:
                dropPadMode.range_values = 0
            elif dropPadMode.range_values == 3:
                dropPadMode.range_values = 2
            E = biasOptional.shape[0]
            Num_Rows = out.shape[0]
            H = expandedX.shape[1]
            K = scalesOptional.shape[1]
            if K > E:
                K = E
                scalesOptional.shape[1] = K # K: 为从总的专家E中选出K个专家;  E: expert num，即专家数，E需要大于等于K; 
            expandedX.shape = [Num_Rows * K, H] # drop less 场景shape为（NUM\_ROWS \* K, H）
            expandedRowIdx.shape = [Num_Rows * K]
            expandedRowIdx.range_values = [0, Num_Rows * K - 1] # dropPadMode参数值为0、2时，Tensor中的值取值范围是[0,NUM\_ROWS \* K-1]; 
            x1Optional.shape = [Num_Rows, H] # x1Optional shape要求与out的shape一致
            x1Optional.dtype = expandedX.dtype # x1Optional 数据类型要求与expandedX一致
            x2Optional.shape = [Num_Rows, H] # x2Optional shape要求与out的shape一致
            x2Optional.dtype = expandedX.dtype # x2Optional 数据类型要求与expandedX一致
            biasOptional.shape = [E, H] # 限制：其shape支持（E，H）
            biasOptional.dtype = expandedX.dtype # 数据类型要求与expandedX一致
            scalesOptional.shape = [Num_Rows, K] # 限制：其shape支持（NUM\_ROWS，K）
            scalesOptional.dtype = expandedX.dtype # 数据类型要求与expandedX一致
            expertIdxOptional.shape = [Num_Rows, K] # 限制：其shape支持（NUM\_ROWS，K）
            expertIdxOptional.range_values = [0, E -1] # 限制：Tensor中的值取值范围是[0, E-1]
            out.dtype = expandedX.dtype
        # print(case_config)
        return case_config